import os
import sys
import yaml
import pandas as pd
import numpy as np
from tqdm import tqdm

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from src.data_prep.data_loader import load_hardware_data, process_hardware_features, split_train_test
from src.features.feature_engineer import HardwareFeatureEngineer
from src.models.model_trainer import train_hardware_model


def generate_mock_hardware_data(num_samples: int = 1000) -> pd.DataFrame:
    """
    生成模拟硬件数据用于演示
    """
    print(f"正在生成{num_samples}条模拟硬件数据...")
    
    np.random.seed(42)
    data = []
    
    for i in tqdm(range(num_samples)):
        device_id = f"ESL-{np.random.randint(100, 300)}-{str(i).zfill(3)}"
        days_in_use = np.random.randint(10, 365)
        avg_voltage = np.random.uniform(2.8, 3.6)
        
        # 电压下降率与电池故障相关
        base_drop_rate = 0.005
        voltage_factor = (3.6 - avg_voltage) * 0.02  # 电压越低，下降率越高
        age_factor = days_in_use / 365 * 0.01  # 使用时间越长，下降率越高
        random_factor = np.random.normal(0, 0.002)
        voltage_drop_rate = max(0.001, base_drop_rate + voltage_factor + age_factor + random_factor)
        
        signal_strength_avg = np.random.uniform(-100, -70)
        screen_refresh_count = np.random.randint(500, 5000)
        
        # 温度统计
        temp_base = np.random.uniform(20, 30)
        temp_max = temp_base + np.random.uniform(5, 20)
        temp_min = temp_base - np.random.uniform(10, 30)
        temp_var = np.random.uniform(5, 25)
        
        # 固件版本和位置
        firmware_version = np.random.choice(['V1.2.3', 'V1.3.0', 'V2.0.1'])
        location_cluster = np.random.choice(['A12', 'B05', 'C08'])
        
        # 生成标签：电池故障概率
        # 基于电压、下降率和使用天数的复合因素
        battery_failure_score = (
            (3.2 - avg_voltage) * 0.5 +  # 低电压增加风险
            voltage_drop_rate * 20 +     # 高下降率增加风险
            days_in_use / 365 * 0.3      # 使用时间增加风险
        )
        # 添加随机因素
        battery_failure_score += np.random.normal(0, 0.1)
        # 转换为概率
        battery_failure_prob = 1 / (1 + np.exp(-(battery_failure_score - 0.5)))  # sigmoid函数
        
        # 二分类标签：未来7天是否会出现电池故障
        battery_failure_in_7d = 1 if battery_failure_prob > 0.7 else 0
        
        # 屏幕故障概率
        screen_failure_prob = (
            screen_refresh_count / 5000 * 0.4 +  # 刷新次数增加风险
            temp_var / 25 * 0.3 +               # 温度波动增加风险
            (35 - avg_voltage) * 0.1            # 电压不稳定影响屏幕
        )
        # 添加随机因素
        screen_failure_prob += np.random.normal(0, 0.1)
        screen_failure_prob = min(0.95, max(0.01, screen_failure_prob))
        
        data.append({
            'device_id': device_id,
            'days_in_use': days_in_use,
            'avg_voltage': round(avg_voltage, 2),
            'voltage_drop_rate': round(voltage_drop_rate, 3),
            'signal_strength_avg': round(signal_strength_avg, 2),
            'screen_refresh_count': screen_refresh_count,
            'temp_max': round(temp_max, 2),
            'temp_min': round(temp_min, 2),
            'temp_var': round(temp_var, 2),
            'firmware_version': firmware_version,
            'location_cluster': location_cluster,
            'battery_failure_in_7d': battery_failure_in_7d,
            'screen_failure_prob': round(screen_failure_prob, 3)
        })
    
    return pd.DataFrame(data)


def main():
    """
    主函数
    """
    # 加载配置
    config_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
                             'config', 'config.yaml')
    with open(config_path, 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)
    
    # 创建数据目录
    os.makedirs(config['data']['sample_data_path'], exist_ok=True)
    
    # 生成模拟数据
    hardware_data_path = os.path.join(config['data']['sample_data_path'], 'hardware_sample_data.csv')
    if not os.path.exists(hardware_data_path):
        df = generate_mock_hardware_data(2000)
        df.to_csv(hardware_data_path, index=False)
        print(f"模拟数据已保存至: {hardware_data_path}")
    else:
        print(f"使用已有的模拟数据: {hardware_data_path}")
        df = pd.read_csv(hardware_data_path)
    
    # 处理数据
    print("开始处理数据...")
    df = process_hardware_features(df)
    
    # 特征工程
    print("开始特征工程...")
    feature_engineer = HardwareFeatureEngineer()
    df = feature_engineer.build_features(df)
    
    # 准备特征和标签
    # 排除ID列和标签列
    exclude_cols = ['device_id', 'battery_failure_in_7d', 'screen_failure_prob']
    feature_cols = [col for col in df.columns if col not in exclude_cols]
    
    X = df[feature_cols]
    y_battery = df['battery_failure_in_7d']
    y_screen = df['screen_failure_prob']
    
    # 划分训练集和测试集
    print("划分训练集和测试集...")
    X_train, X_test, y_train_battery, y_test_battery = split_train_test(X, y_battery)
    _, _, y_train_screen, y_test_screen = split_train_test(X, y_screen)  # 使用相同的划分
    
    print(f"训练集大小: {len(X_train)}, 测试集大小: {len(X_test)}")
    
    # 训练模型
    print("开始训练硬件故障预测模型...")
    model = train_hardware_model(
        config=config,
        X_train=X_train,
        y_train_battery=y_train_battery,
        y_train_screen=y_train_screen,
        X_test=X_test,
        y_test_battery=y_test_battery,
        y_test_screen=y_test_screen
    )
    
    print("\n模型训练完成！")
    print(f"电池模型保存路径: {config['models']['hardware_model_path'].replace('.txt', '_battery.txt')}")
    print(f"屏幕模型保存路径: {config['models']['hardware_model_path'].replace('.txt', '_screen.txt')}")


if __name__ == "__main__":
    main()
